package com.wispx.gateway.core.filter.flow.algorithm;

import com.wispx.gateway.core.context.GatewayContext;
import com.wispx.gateway.core.filter.flow.RateLimiter;

/**
 * 令牌桶算法
 */
public class TokenBucketRateLimter implements RateLimiter {

    /**
     * 上一次访问的时间戳，单位 ms
     */
    private long lastUpdateTime = System.currentTimeMillis();

    /**
     * 桶中剩余的令牌数
     */
    private int currentTokens;

    /**
     * 桶的容量（最大请求次数）
     */
    private int capacity;

    /**
     * 生成 Token 的速率，单位 个/s  5ms 1个 -> 1000ms 200个
     */
    private int rate;

    public TokenBucketRateLimter(int capacity, int rate) {
        this.capacity = capacity;
        this.rate = rate;
        this.currentTokens = capacity;
    }

    @Override
    public void tryConsume(GatewayContext context) {
        if (takeToken()) {
            context.doFilter();
        } else {
            throw new RuntimeException("流量到达限流上限，该请求将被丢弃");
        }
    }

    private synchronized boolean takeToken() {
        adjust();
        if (currentTokens > 0) {
            currentTokens--;
            return true;
        }
        return false;
    }

    private void adjust() {
        if (currentTokens == capacity) { // 桶满了
            return;
        }
        long now = System.currentTimeMillis();
        long timeDiff = now - lastUpdateTime;
        int newTokens = (int) (currentTokens + ((timeDiff + 1000 - 1) / 1000) * rate);
        if (newTokens == 0) {
            return;
        }
        currentTokens = Math.min(capacity, newTokens);
        lastUpdateTime = now;
    }
}
